import plotly.graph_objects as go
import numpy as np

# 定义两个复数
z1 = 2 * (np.cos(np.pi / 4) + 1j * np.sin(np.pi / 4))
z2 = 1 * (np.cos(np.pi / 6) + 1j * np.sin(np.pi / 6))

# 计算复数乘积
z_product = z1 * z2

# 获取实部、虚部和模长
x1, y1 = z1.real, z1.imag
r1 = np.abs(z1)
x2, y2 = z2.real, z2.imag
r2 = np.abs(z2)
x_product, y_product = z_product.real, z_product.imag
r_product = np.abs(z_product)

# 定义三棱锥的顶点
vertices = [
    [0, 0, 0],
    [x1, y1, r1],
    [x2, y2, r2],
    [x_product, y_product, r_product]
]

# 定义三棱锥的面
faces = [
    [0, 1, 2],
    [0, 1, 3],
    [0, 2, 3],
    [1, 2, 3]
]

# 提取顶点的坐标
x = [v[0] for v in vertices]
y = [v[1] for v in vertices]
z = [v[2] for v in vertices]

# 创建三棱锥的三角面
i, j, k = [], [], []
for face in faces:
    i.append(face[0])
    j.append(face[1])
    k.append(face[2])

# 创建 3D 图形
fig = go.Figure(data=[
    go.Mesh3d(
        x=x, y=y, z=z,
        i=i, j=j, k=k,
        opacity=0.5,
        color='cyan'
    )
])

# 绘制代表 z1 的边并标注，设置颜色为红色，设置图例名称
fig.add_trace(
    go.Scatter3d(
        x=[0, x1], y=[0, y1], z=[0, r1],
        mode='lines+text',
        text=['$z_1$'],
        textposition='top center',
        line=dict(color='red', width=3),
        name='$z_1$'
    )
)

# 绘制代表 z2 的边并标注，设置颜色为绿色，设置图例名称
fig.add_trace(
    go.Scatter3d(
        x=[0, x2], y=[0, y2], z=[0, r2],
        mode='lines+text',
        text=['$z_2$'],
        textposition='top center',
        line=dict(color='green', width=3),
        name='$z_2$'
    )
)

# 绘制代表 z1z2 的边并标注，设置颜色为蓝色，设置图例名称
fig.add_trace(
    go.Scatter3d(
        x=[0, x_product], y=[0, y_product], z=[0, r_product],
        mode='lines+text',
        text=['$z_1z_2$'],
        textposition='top center',
        line=dict(color='blue', width=3),
        name='$z_1z_2$'
    )
)

# 设置图形布局
fig.update_layout(
    scene=dict(
        xaxis_title='Real Axis',
        yaxis_title='Imaginary Axis',
        zaxis_title='Magnitude'
    )
)

# 显示图形
fig.show()
